{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- bert encoder "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encoder(inputs):\n",
    "    Q = tf.keras.layers.Dense(64, name = 'Queries')(inputs)\n",
    "    K = tf.keras.layers.Dense(64, name = 'Key')(inputs)\n",
    "    V = tf.keras.layers.Dense(64, name = 'Values')(inputs)\n",
    "    score = tf.keras.layers.Softmax(tf.keras.layers.LayerNormalization(tf.multiply(Q,K)))\n",
    "    Z = tf.multiply(score,V)\n",
    "    encoder_output = tf.keras.layers.Dense(128)(tf.keras.layers.LayerNormalization(Z))\n",
    "    return encoder_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- bert position information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionEmbedding(tf.keras.layers.Layer):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_dim,\n",
    "        output_dim,\n",
    "        merge_mode='add',\n",
    "        hierarchical=None,\n",
    "        embeddings_initializer='zeros',\n",
    "        custom_position_ids=False,\n",
    "        **kwargs\n",
    "    ):\n",
    "        super(PositionEmbedding, self).__init__(**kwargs)\n",
    "        self.input_dim = input_dim\n",
    "        self.output_dim = output_dim\n",
    "        self.merge_mode = merge_mode\n",
    "        self.hierarchical = hierarchical\n",
    "        self.embeddings_initializer = initializers.get(embeddings_initializer)\n",
    "        self.custom_position_ids = custom_position_ids\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        super(PositionEmbedding, self).build(input_shape)\n",
    "        self.embeddings = self.add_weight(\n",
    "            name='embeddings',\n",
    "            shape=(self.input_dim, self.output_dim),\n",
    "            initializer=self.embeddings_initializer\n",
    "        )\n",
    "        \n",
    "    def call(self, inputs):\n",
    "        \"\"\"如果custom_position_ids，那么第二个输入为自定义的位置id\n",
    "        \"\"\"\n",
    "        if self.custom_position_ids:\n",
    "            inputs, position_ids = inputs\n",
    "            if 'int' not in K.dtype(position_ids):\n",
    "                position_ids = K.cast(position_ids, 'int32')\n",
    "        else:\n",
    "            input_shape = K.shape(inputs)\n",
    "            batch_size, seq_len = input_shape[0], input_shape[1]\n",
    "            position_ids = K.arange(0, seq_len, dtype='int32')[None]\n",
    "\n",
    "        if self.hierarchical:\n",
    "            alpha = 0.4 if self.hierarchical is True else self.hierarchical\n",
    "            embeddings = self.embeddings - alpha * self.embeddings[:1]\n",
    "            embeddings = embeddings / (1 - alpha)\n",
    "            embeddings_x = K.gather(embeddings, position_ids // self.input_dim)\n",
    "            embeddings_y = K.gather(embeddings, position_ids % self.input_dim)\n",
    "            embeddings = alpha * embeddings_x + (1 - alpha) * embeddings_y\n",
    "        else:\n",
    "            if self.custom_position_ids:\n",
    "                embeddings = K.gather(self.embeddings, position_ids)\n",
    "            else:\n",
    "                embeddings = self.embeddings[None, :seq_len]\n",
    "\n",
    "        if self.merge_mode == 'add':\n",
    "            return inputs + embeddings\n",
    "        elif self.merge_mode == 'mul':\n",
    "            return inputs * (embeddings + 1.0)\n",
    "        elif self.merge_mode == 'zero':\n",
    "            return embeddings\n",
    "        else:\n",
    "            if not self.custom_position_ids:\n",
    "                embeddings = K.tile(embeddings, [batch_size, 1, 1])\n",
    "            return K.concatenate([inputs, embeddings])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
